import json
from functools import partial
from operator import itemgetter

import numpy as np
import requests
from fastapi import FastAPI
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
origins = ["*"]


def get_embedding(text, model="vicuna-13b-v1.5-16k"):
    text = text.replace("\n", " ")
    payload = {
        "model": "vicuna-13b-v1.5-16k",
        "input": text
    }
   # return client.embeddings.create(input = [text], model=model).data
    data = requests.post("http://113.31.110.212:7001/v1/embeddings",data=json.dumps(payload)).json()['data'][0]['embedding']
    return data

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

with open('./embeds.json','r') as f:
    data = json.load(f)

def cosine_similarity(vector_a, vector_b):
    dot_product = np.dot(vector_a, vector_b)
    norm_a = np.linalg.norm(vector_a)
    norm_b = np.linalg.norm(vector_b)
    similarity = dot_product / (norm_a * norm_b)
    return similarity

@app.get("/sim/{text}")
def get_sim(text:str):
     target_embed = np.array(get_embedding(text))
     cos_fn = partial(cosine_similarity,target_embed)
     all = []
     for k,v in data.items():
         v1 = np.array(v[0])
         v2 = np.array(v[1])
         v3 = np.array(v[2])
         vv = [v1,v2,v3]
         res = list(map(cos_fn,vv))
         all.append((k,np.mean(res)))
     _all = sorted(all,key=itemgetter(1),reverse=True)
     return {"code":200,"data":_all[0:10]}



if __name__ == '__main__':
    uvicorn.run(app,host="0.0.0.0",port=5000)



